package com.childenglish.filter;

import org.owasp.html.HtmlPolicyBuilder;
import org.owasp.html.PolicyFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.IOException;
import java.util.regex.Pattern;

/**
 * XSS防护过滤器
 * 过滤请求参数中的XSS攻击代码
 */
public class XSSFilter implements Filter {

    private static final Logger logger = LoggerFactory.getLogger(XSSFilter.class);

    // XSS攻击模式
    private static final Pattern[] XSS_PATTERNS = {
        Pattern.compile("<script[^>]*>.*?</script>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL),
        Pattern.compile("<iframe[^>]*>.*?</iframe>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL),
        Pattern.compile("javascript:", Pattern.CASE_INSENSITIVE),
        Pattern.compile("on\\w+\\s*=", Pattern.CASE_INSENSITIVE),
        Pattern.compile("<img[^>]*src[^>]*=.*javascript:", Pattern.CASE_INSENSITIVE),
        Pattern.compile("<svg[^>]*onload", Pattern.CASE_INSENSITIVE),
        Pattern.compile("expression\\s*\\(", Pattern.CASE_INSENSITIVE),
        Pattern.compile("vbscript:", Pattern.CASE_INSENSITIVE),
        Pattern.compile("<style[^>]*>.*?</style>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL)
    };

    // OWASP HTML Sanitizer策略
    private static final PolicyFactory POLICY = new HtmlPolicyBuilder()
        .allowElements("p", "br", "strong", "em", "u", "h1", "h2", "h3", "h4", "h5", "h6")
        .allowAttributes("class").onElements("p", "div", "span")
        .toFactory();

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        XSSRequestWrapper wrappedRequest = new XSSRequestWrapper((HttpServletRequest) request);
        chain.doFilter(wrappedRequest, response);
    }

    /**
     * XSS请求包装器
     */
    private static class XSSRequestWrapper extends HttpServletRequestWrapper {

        public XSSRequestWrapper(HttpServletRequest request) {
            super(request);
        }

        @Override
        public String[] getParameterValues(String parameter) {
            String[] values = super.getParameterValues(parameter);
            if (values == null) {
                return null;
            }
            int count = values.length;
            String[] encodedValues = new String[count];
            for (int i = 0; i < count; i++) {
                encodedValues[i] = cleanXSS(values[i]);
            }
            return encodedValues;
        }

        @Override
        public String getParameter(String parameter) {
            String value = super.getParameter(parameter);
            return cleanXSS(value);
        }

        @Override
        public String getHeader(String name) {
            String value = super.getHeader(name);
            return cleanXSS(value);
        }

        /**
         * 清理XSS攻击代码
         */
        private String cleanXSS(String value) {
            if (value == null) {
                return null;
            }

            // 检测XSS攻击
            for (Pattern pattern : XSS_PATTERNS) {
                if (pattern.matcher(value).find()) {
                    logger.warn("检测到XSS攻击尝试: {}", value);
                    // 使用OWASP HTML Sanitizer清理
                    value = POLICY.sanitize(value);
                    // 如果清理后仍有问题，转义HTML
                    if (value == null || value.isEmpty()) {
                        value = escapeHtml(value);
                    }
                }
            }

            // 转义HTML特殊字符
            value = escapeHtml(value);

            return value;
        }

        /**
         * 转义HTML特殊字符
         */
        private String escapeHtml(String value) {
            if (value == null) {
                return null;
            }
            return value
                .replace("&", "&amp;")
                .replace("<", "&lt;")
                .replace(">", "&gt;")
                .replace("\"", "&quot;")
                .replace("'", "&#x27;")
                .replace("/", "&#x2F;");
        }
    }
}

